package com.example.calcite.mysql;

import org.apache.calcite.schema.Table;
import org.apache.calcite.schema.impl.AbstractSchema;

import java.sql.Connection;
import java.sql.DriverManager;
import java.sql.ResultSet;
import java.sql.SQLException;
import java.util.HashMap;
import java.util.Map;
import java.util.Properties;
import java.util.regex.Matcher;
import java.util.regex.Pattern;

public class MySQLSchema extends AbstractSchema {
    private final String jdbcUrl;
    private final String schema;
    private final Properties properties;
    private final Map<String, Table> tableMap = new HashMap<>();

    public MySQLSchema(String jdbcUrl, String username, String password) {
        this.jdbcUrl = jdbcUrl;
        this.schema = extractSchemaFromUrl(jdbcUrl);
        this.properties = new Properties();
        this.properties.setProperty("user", username);
        this.properties.setProperty("password", password);
        loadTables();
    }

    private void loadTables() {
        try (Connection connection = DriverManager.getConnection(jdbcUrl, properties)) {
            // 获取所有表名
            ResultSet tables = connection.getMetaData().getTables(schema, null, null, new String[]{"TABLE"});
            while (tables.next()) {
                String tableName = tables.getString("TABLE_NAME");
                tableMap.put(tableName, new MySQLTable(jdbcUrl, properties, tableName));
            }
        } catch (SQLException e) {
            throw new RuntimeException("Failed to load MySQL tables", e);
        }
    }

    @Override
    protected Map<String, Table> getTableMap() {
        return tableMap;
    }

    private String extractSchemaFromUrl(String jdbcUrl) {
        if (jdbcUrl == null || jdbcUrl.isEmpty()) {
            throw new IllegalArgumentException("JDBC URL cannot be null or empty");
        }

        String regex = "jdbc:mysql://[^/]+/([^?]+).*";
        Pattern pattern = Pattern.compile(regex);
        Matcher matcher = pattern.matcher(jdbcUrl);

        if (matcher.matches()) {
            String schema = matcher.group(1);
            if (schema == null || schema.isEmpty()) {
                throw new IllegalArgumentException("No database name found in JDBC URL");
            }
            return schema;
        }

        throw new IllegalArgumentException("Invalid MySQL JDBC URL format");
    }
} 